import torch
import torchvision
import torch.nn.functional as F

import numpy as np
import pandas as pd
from PIL import Image


def classify(img_path):
    import matplotlib.pyplot as plt
    plt.show()

    # 有 GPU 就用 GPU，没有就用 CPU
    device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')

    # windows操作系统
    plt.rcParams['font.sans-serif'] = ['SimHei']  # 用来正常显示中文标签
    plt.rcParams['axes.unicode_minus'] = False  # 用来正常显示负号

    from PIL import Image, ImageFont, ImageDraw
    # 导入中文字体，指定字号
    font = ImageFont.truetype('SimHei.ttf', 32)

    # 载入类别
    idx_to_labels = np.load('idx_to_labels.npy', allow_pickle=True).item()

    # 导入训练好的模型
    model = torch.load('checkpoint/melon_pytorch_C1.pth')
    model = model.eval().to(device)

    # 预处理
    from torchvision import transforms
    # 测试集图像预处理-RCTN：缩放、裁剪、转 Tensor、归一化
    test_transform = transforms.Compose([transforms.Resize(256),
                                         transforms.CenterCrop(224),
                                         transforms.ToTensor(),
                                         transforms.Normalize(
                                             mean=[0.485, 0.456, 0.406],
                                             std=[0.229, 0.224, 0.225])
                                         ])

    # img_path = 'test_img/img_1.png'
    img_pil = Image.open(img_path)
    np.array(img_pil).shape

    input_img = test_transform(img_pil)  # 预处理
    # input_img.shape
    input_img = input_img.unsqueeze(0).to(device)
    # input_img.shape
    # 执行前向预测，得到所有类别的 logit 预测分数
    pred_logits = model(input_img)
    pred_softmax = F.softmax(pred_logits, dim=1)  # 对 logit 分数做 softmax 运算

    plt.figure(figsize=(22, 10))

    x = idx_to_labels.values()
    y = pred_softmax.cpu().detach().numpy()[0] * 100
    width = 0.45  # 柱状图宽度

    ax = plt.bar(x, y, width)
    y_max = np.argmax(y)
    # 得到识别的结果
    result = list(x)[y_max]
    print(result)

    plt.bar_label(ax, fmt='%.2f', fontsize=15)  # 置信度数值
    plt.tick_params(labelsize=15)  # 设置坐标文字大小

    plt.title(img_path, fontsize=30)
    plt.xticks(rotation=45)  # 横轴文字旋转
    plt.xlabel('类别', fontsize=20)
    plt.ylabel('置信度', fontsize=20)
    plt.show()

    return result